package com.github.kezhenxu94.data.source;

import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.clients.consumer.ConsumerRecords;
import org.apache.kafka.clients.consumer.KafkaConsumer;

import java.util.Collection;
import java.util.Properties;
import java.util.concurrent.*;

/**
 * Created by kezhenxu on 1/5/17.
 *
 * @author kezhenxu (kezhenxu94@163.com)
 */
public class KafkaSource<K, V> implements Source<V> {

    protected Properties properties;
    protected KafkaConsumer<K, V> kafkaConsumer;
    protected BlockingQueue<V> queue;
    protected Collection<String> topics;
    protected ExecutorService executor;

    private volatile boolean isRunning;

    @Override
    public void setup() throws Exception {
        executor = Executors.newSingleThreadExecutor();
        properties = new Properties();
        properties.load(getClass().getResourceAsStream("/kafka-source.properties"));
        kafkaConsumer = new KafkaConsumer<>(properties);
        kafkaConsumer.subscribe(topics);
        queue = new LinkedBlockingQueue<>(200);
    }

    @Override
    public void start() throws Exception {
        isRunning = true;
        executor.execute(() -> {
            while (isRunning()) {
                ConsumerRecords<K, V> records = kafkaConsumer.poll(200);
                for (ConsumerRecord<K, V> record : records) {
                    queue.offer(record.value());
                }
            }
        });
    }

    @Override
    public void close() throws Exception {
        isRunning = false;
        while (isRunning()) {
            TimeUnit.MICROSECONDS.sleep(200);
        }
        executor.awaitTermination(500L, TimeUnit.MILLISECONDS);
        if (kafkaConsumer != null) {
            kafkaConsumer.close();
            kafkaConsumer = null;
        }
    }

    @Override
    public boolean isRunning() {
        return isRunning || !queue.isEmpty();
    }

    @Override
    public V take() throws Exception {
        return isRunning() ? queue.take() : null;
    }

    public Collection<String> getTopics() {
        return topics;
    }

    public void setTopics(Collection<String> topics) {
        this.topics = topics;
    }
}
